package com.example.androidaudiovideodemo.utils

import android.media.MediaCodec
import android.media.MediaCodecInfo
import android.media.MediaFormat
import android.util.Log
import android.view.Surface
import java.io.*
import java.nio.ByteBuffer
import java.util.concurrent.ArrayBlockingQueue
import kotlin.concurrent.thread

/**
 *
 * author: renbing
 *
 * date: 2022/3/16
 *
 * des:avc编码
 */
object AvcEncodeUtil {
    private val TAG = "AvcEncoder"

    val mYuvQueue = ArrayBlockingQueue<ByteArray?>(10)

    private var mMediaCodec: MediaCodec? = null
    private var mWidth = 0
    private var mHeight = 0
    private var mFrameRate = 0

    private var mOutFile: File? = null
    private var mOutputStream: BufferedOutputStream? = null
    private var mIsCamera = false
    private var mConfigByte: ByteArray? = null

    private var mIsRunning = false
    private var mStopFlag = false


    fun initData(width: Int,height: Int,frameRate: Int, outFile: File?,isCamera: Boolean){
        mWidth = width
        mHeight = height
        mFrameRate = frameRate
        mOutFile = outFile
        mIsCamera = isCamera

        val mediaFormat = MediaFormat.createVideoFormat(MediaFormat.MIMETYPE_VIDEO_AVC,height, width)
        mediaFormat.run {
            setInteger(MediaFormat.KEY_COLOR_FORMAT,MediaCodecInfo.CodecCapabilities.COLOR_FormatYUV420Flexible)
            setInteger(MediaFormat.KEY_BIT_RATE,width*height*5)//设置码率
            setInteger(MediaFormat.KEY_FRAME_RATE,mFrameRate)//设置帧率
            setInteger(MediaFormat.KEY_I_FRAME_INTERVAL,1)//关键帧间隔时间（I帧）
        }

        try {
            mMediaCodec = MediaCodec.createEncoderByType(MediaFormat.MIMETYPE_VIDEO_AVC)
        }catch (e: IOException){
            e.printStackTrace()
        }
        mMediaCodec?.configure(mediaFormat,null,null,MediaCodec.CONFIGURE_FLAG_ENCODE)
        mMediaCodec?.start()

        createFile()
    }

    private fun createFile(){
        try {
            mOutputStream = BufferedOutputStream(FileOutputStream(mOutFile))
        }catch (e: FileNotFoundException){
            e.printStackTrace()
        }
    }

    fun putYuvData(buffer: ByteArray?){
        if (mYuvQueue.size >= 10){
            mYuvQueue.poll()
        }
        mYuvQueue.add(buffer)
    }

    fun stopThread(){
        if (!mIsRunning) return
        mIsRunning = false
        try {
            stopEncode()
            mOutputStream?.flush()
            mOutputStream?.close()
            mOutputStream = null
        }catch (e: Exception){
            e.printStackTrace()
        }
    }

    private fun stopEncode(){
        mMediaCodec?.stop()
        mMediaCodec?.release()
        mMediaCodec = null
    }

    /**
     * 开启一个子线程开始解码
     */
    fun startThread(){
        thread {
            mIsRunning = true
            var input: ByteArray? = null
            var pts = 0L
            var generateIndex = 0L
            while (mIsRunning){
                if (mYuvQueue.size > 0){
                    input = mYuvQueue.poll()
                    //YUV 1个像素点占用1个Y分量 4个像素公用一个U分量和1个V分量 所以空间大小为mWidth*mHeight*(1+1/4+1/4) ==>mWidth*mHeight*3/2
                    val yuv420sp = ByteArray(mWidth* mHeight*3/2)
                    if (mIsCamera){
                        NV21ToNV12(input,yuv420sp, mWidth, mHeight)
                    }else{
                        YV12ToNV12(input,yuv420sp, mWidth, mHeight)
                    }
                    input = yuv420sp
                }

                if (input != null){
                    try {
                        //从Input缓冲区获取可用的空的InputBuffer索引，传入-1表示一直等待直到返回可用的索引
                        val inputBufferIndex = mMediaCodec?.dequeueInputBuffer(-1)
                        Log.e("TAG","inputBufferIndex:$inputBufferIndex")
                        if (inputBufferIndex != null && inputBufferIndex >= 0){
                            //计算出演示时间戳
                            pts = computePresentationTime(generateIndex)
                            //根据索引获取可用的InputBuffer
                            val inputBuffer = mMediaCodec?.getInputBuffer(inputBufferIndex)
                            //清除InputBuffer中之前的旧数据
                            inputBuffer?.clear()
                            //放入待解码的新数据
                            inputBuffer?.put(input)
                            //放回Input缓冲区
                            mMediaCodec?.queueInputBuffer(inputBufferIndex,0,input.size,pts,0)
                            generateIndex +=1
                        }

                        val bufferInfo = MediaCodec.BufferInfo()
                        //从output缓冲区获取装填有解码后的数据的OutputBuffer索引
                        var outputBufferIndex = mMediaCodec?.dequeueOutputBuffer(bufferInfo,-1)
                        while (outputBufferIndex != null && outputBufferIndex >= 0){
                            //获取装填有解码数据的outputbuffer
                            val outputBuffer = mMediaCodec?.getOutputBuffer(outputBufferIndex)
                            val outData = ByteArray(bufferInfo.size)
                            //转移解码数据到ByteArray对象中
                            outputBuffer?.get(outData)
                            when(bufferInfo.flags){
                                MediaCodec.BUFFER_FLAG_CODEC_CONFIG ->{
                                    mConfigByte = ByteArray(bufferInfo.size)
                                    mConfigByte = outData
                                }
                                MediaCodec.BUFFER_FLAG_KEY_FRAME ->{
                                    val keyFrame = ByteArray(bufferInfo.size+ mConfigByte?.size!!)
                                    System.arraycopy(mConfigByte?:ByteArray(bufferInfo.size),0,keyFrame,0, mConfigByte?.size!!)
                                    System.arraycopy(outData,0,keyFrame, mConfigByte?.size!!, outData.size)
                                    try {
                                        mOutputStream?.write(keyFrame,0,keyFrame.size)
                                    }catch (e: IOException){
                                        e.printStackTrace()
                                    }
                                }
                                else ->{
                                    try {
                                        //输出到文件流
                                        mOutputStream?.write(outData,0,outData.size)
                                    }catch (e: IOException){
                                        e.printStackTrace()
                                    }
                                }
                            }
                            //释放OutputBuffer
                            mMediaCodec?.releaseOutputBuffer(outputBufferIndex,false)
                            outputBufferIndex = mMediaCodec?.dequeueOutputBuffer(bufferInfo,0)
                        }
                    }catch (e: Exception){
                        e.printStackTrace()
                    }
                }else{
                    try {
                        Thread.sleep(500)
                    }catch (e: InterruptedException){
                        e.printStackTrace()
                    }
                }
            }
        }
    }

    private fun NV21ToNV12(nv21: ByteArray?,nv12: ByteArray?,width: Int,height: Int){
        if (nv21 == null || nv12 == null) return
        val frameSize = mWidth * mHeight
        System.arraycopy(nv21,0,nv12,0,frameSize)
        for (i in 0 until frameSize){
            nv12[i] = nv21[i]
        }

        for (j in 0 until frameSize/2 step 2){
            nv12[frameSize+j-1] = nv21[j+frameSize]
        }

        for (k in 0 until frameSize/2 step 2){
            nv12[frameSize+k] = nv21[k+frameSize-1]
        }
    }

    private fun YV12ToNV12(yv12: ByteArray?,nv12: ByteArray?,width: Int,height: Int){
        if (yv12 == null || nv12 == null) return
        val lenY = width* height
        val lenU = lenY/4
        System.arraycopy(yv12,0,nv12,0,width*height)
        for (i in 0 until lenU){
            nv12[lenY+2*i] = yv12[lenY+i]
            nv12[lenY+2*i+1] = yv12[lenY+lenU+i]
        }
    }

    private fun computePresentationTime(frameIndex: Long): Long{
        return 132+frameIndex*1000000/mFrameRate
    }

    fun encodeH264(data : ByteArray,encodeData:(ByteArray) -> Unit){
        val bytes = NV21ToNV12(data,mWidth,mHeight)
        val nv12 = rotate90(bytes,mWidth,mHeight)
        try {
            val inputBufferIndex = mMediaCodec?.dequeueInputBuffer(-1)?:-1
            if (inputBufferIndex >= 0){
                val inputBuffer = mMediaCodec?.getInputBuffer(inputBufferIndex)
                inputBuffer?.clear()
                inputBuffer?.put(nv12)
                mMediaCodec?.queueInputBuffer(inputBufferIndex,0,nv12.size,System.nanoTime()/1000,0)
            }
            val bufferInfo = MediaCodec.BufferInfo()
            var outputBufferIndex = mMediaCodec?.dequeueOutputBuffer(bufferInfo,0)?:-1
            while (outputBufferIndex >= 0){
                val outputBuffer = mMediaCodec?.getOutputBuffer(outputBufferIndex)
                val outData = ByteArray(bufferInfo.size)
                outputBuffer?.get(outData)
                encodeData.invoke(outData)
                mOutputStream?.write(outData)
                mMediaCodec?.releaseOutputBuffer(outputBufferIndex,false)
                outputBufferIndex = mMediaCodec?.dequeueOutputBuffer(bufferInfo,0)?:-1
            }
        }catch (e:Exception){
            e.printStackTrace()
        }
    }

    fun decodeH264(width: Int,height: Int,surface: Surface,ins: DataInputStream){
        try {
            val decodeCodec = MediaCodec.createDecoderByType(MediaFormat.MIMETYPE_VIDEO_AVC)
            val decodeFormat = MediaFormat.createVideoFormat(MediaFormat.MIMETYPE_VIDEO_AVC,width, height)
            val headerSPS = byteArrayOf(0, 0, 0, 1, 67, 66, 0, 42,149.toByte(), 168.toByte(), 30, 0, 137.toByte(), 249.toByte(), 102,224.toByte(), 32, 32, 32, 64)
            val headerPPS = byteArrayOf(0, 0, 0, 1, 68, 206.toByte(), 60, 128.toByte(), 0, 0, 0, 1, 6, 229.toByte(), 1, 151.toByte(), 128.toByte())
            decodeFormat.run {
                setByteBuffer("csd-0",ByteBuffer.wrap(headerSPS))
                setByteBuffer("csd-1",ByteBuffer.wrap(headerPPS))
                setInteger(MediaFormat.KEY_FRAME_RATE,40)
            }
            decodeCodec.configure(decodeFormat,surface,null,0)
            decodeCodec.start()

            val bufferInfo = MediaCodec.BufferInfo()
            val startMs = System.currentTimeMillis()
            val timeoutUs = -1L
            val maker0 = byteArrayOf(0,0,0,1)
            val dummyFrame = byteArrayOf(0x00,0x00,0x01,0x20)
            var streamBuffer: ByteArray? = null
            try {
                streamBuffer = getBytes(ins)
            }catch (e: Exception){
                e.printStackTrace()
            }
            var bytesCnt = 0
            while (!mStopFlag){
                bytesCnt = streamBuffer?.size?:0
                if (bytesCnt == 0){
                    streamBuffer = dummyFrame
                }
                var startIndex = 0
                var remaining = bytesCnt
                while (true){
                    if (remaining == 0 || startIndex >= remaining){
                        break
                    }
                    var nextFrameStart = KMPMatch(maker0,streamBuffer!!,startIndex+2,remaining)
                    if (nextFrameStart == -1){
                        nextFrameStart = remaining
                    }
                    val inputIndex = decodeCodec.dequeueInputBuffer(0)
                    if (inputIndex >= 0){
                        val byteBuffer = decodeCodec.getInputBuffer(inputIndex)
                        byteBuffer?.clear()
                        byteBuffer?.put(streamBuffer,startIndex,nextFrameStart-startIndex)
                        decodeCodec.queueInputBuffer(inputIndex,0,nextFrameStart-startIndex,0,0)
                        startIndex = nextFrameStart
                    }else{
                        continue
                    }

                    var outputIndex = decodeCodec.dequeueOutputBuffer(bufferInfo,0)
                    if (outputIndex >= 0){
                        while (bufferInfo.presentationTimeUs/1000 > System.currentTimeMillis()-startMs){
                            try {
                                Thread.sleep(100)
                            }catch (e: Exception){
                                e.printStackTrace()
                            }
                        }
                        var doRender = bufferInfo.size != 0
                        decodeCodec.releaseOutputBuffer(outputIndex,doRender)
                    }
                }
                mStopFlag = false
            }
        }catch (e: Exception){
            e.printStackTrace()
        }

    }

    private fun getBytes(ins: InputStream): ByteArray?{
        var len = 0
        var size = 1024
        var buf: ByteArray?
        if (ins is ByteArrayInputStream){
            size = ins.available()
            buf = ByteArray(size)
            len = ins.read(buf,0,size)
        }else{
            val bos = ByteArrayOutputStream()
            buf = ByteArray(size)
            while (ins.read(buf,0,size).also { len = it } != -1) //将读取的数据写入到字节输出流
                bos.write(buf,0,len) //将这个流转换成字节数组
            //将这个流转换成字节数组
            buf = bos.toByteArray()
        }
        return buf
    }

    private fun KMPMatch(pattern: ByteArray,bytes: ByteArray,start: Int,remain: Int): Int{
        try {
            Thread.sleep(30)
        }catch (e: Exception){
            e.printStackTrace()
        }
        val lsp = computeLspTable(pattern)
        var j = 0
        for (i in start until remain){
            while (j > 0 && bytes[i] != pattern[j]){
                j = lsp[j-1]
            }
            if (bytes[i] == pattern[j]){
                j++
                if (j == pattern.size){
                    return i-j+1
                }
            }
        }
        return -1
    }

    private fun computeLspTable(pattern: ByteArray): IntArray{
        val lsp = IntArray(pattern.size)
        for (i in pattern.indices){
            if (i == 0){
                lsp[i] = 0
            }else{
                var j = lsp[i-1]
                while (j > 0 && pattern[i] != pattern[j]){
                    j = lsp[j-1]
                }
                if (pattern[i] == pattern[j]){
                    j++
                }
                lsp[i] = j
            }
        }
        return lsp
    }

    fun NV21ToNV12(nv21: ByteArray,width: Int,height: Int): ByteArray{
        val nv12 = ByteArray(width*height*3/2)
        val frameSize = width*height
        System.arraycopy(nv21,0,nv12,0,frameSize)
        for (i in 0 until frameSize){
            nv12[i] = nv21[i]
        }
        for (j in 0 until frameSize/2 step 2){
            nv12[frameSize+j-1] = nv21[j+frameSize]
        }
        for (j in 0 until frameSize/2 step 2){
            nv12[frameSize+j] = nv21[j+frameSize-1]
        }
        return nv12
    }

    fun rotate90(data: ByteArray,previousWidth: Int,previousHeight: Int): ByteArray{
        val yuv = ByteArray(previousWidth*previousHeight*3/2)
        var i = 0
        for (x in 0 until previousWidth){
            for (y in previousHeight-1 downTo 0){
                yuv[i++] = data[y*previousWidth+x]
            }
        }
        i = previousWidth*previousHeight*3/2-1
        for (x in previousWidth-1 downTo 0 step 2){
            for (y in 0 until previousHeight/2){
                yuv[i--] = data[(previousWidth*previousHeight)+(y*previousWidth)+x]
                yuv[i--] = data[(previousWidth*previousHeight)+(y*previousWidth)+(x-1)]
            }
        }
        return yuv
    }

}